# 版本: 1.1.0
# 作者: 陈振玺
# 功能: CSV/Excel矩阵过滤工具

import pandas as pd
import numpy as np
import argparse
import re
from pathlib import Path

class MatrixFilter:
    def __init__(self, input_file, output_file):
        self.input_file = Path(input_file)
        self.output_file = Path(output_file)
        self.data = None
        self.filters = []
    
    def load_data(self):
        """加载CSV或Excel数据"""
        try:
            if self.input_file.suffix.lower() == '.csv':
                self.data = pd.read_csv(self.input_file)
            elif self.input_file.suffix.lower() in ['.xls', '.xlsx']:
                self.data = pd.read_excel(self.input_file)
            else:
                print(f"错误: 不支持的文件格式 {self.input_file.suffix}")
                return False
            
            print(f"成功加载数据: {len(self.data)} 行, {len(self.data.columns)} 列")
            print(f"可用列名: {list(self.data.columns)}")
            return True
            
        except Exception as e:
            print(f"加载数据失败: {e}")
            return False
    
    def add_filter(self, column, condition):
        """添加过滤条件"""
        if column not in self.data.columns:
            print(f"错误: 列 '{column}' 不存在")
            print(f"可用列名: {list(self.data.columns)}")
            return False
        
        self.filters.append({'column': column, 'condition': condition})
        print(f"添加过滤条件: {column} -> {condition}")
        return True
    
    def parse_condition(self, column, condition):
        """解析过滤条件"""
        column_data = self.data[column]
        
        # 检查是否为空值条件
        if self._is_null_condition(condition):
            return self._parse_null_condition(column, condition)
        
        # 检查列的数据类型
        if pd.api.types.is_numeric_dtype(column_data):
            # 数值型列 - 支持比较运算符
            return self._parse_numeric_condition(column, condition)
        else:
            # 字符串型列 - 支持字符串匹配
            return self._parse_string_condition(column, condition)
    
    def _parse_numeric_condition(self, column, condition):
        """解析数值型条件"""
        try:
            # 支持的数值条件格式:
            # >10, >=10, <10, <=10, ==10, !=10
            # 10-20 (范围)
            # 10,20,30 (多个值)
            
            condition = condition.strip()
            
            # 范围条件 (例: 10-20)
            if '-' in condition and not condition.startswith('-'):
                parts = condition.split('-')
                if len(parts) == 2:
                    try:
                        min_val = float(parts[0])
                        max_val = float(parts[1])
                        return (self.data[column] >= min_val) & (self.data[column] <= max_val)
                    except ValueError:
                        pass
            
            # 多个值条件 (例: 10,20,30)
            if ',' in condition:
                try:
                    values = [float(x.strip()) for x in condition.split(',')]
                    return self.data[column].isin(values)
                except ValueError:
                    pass
            
            # 比较运算符条件
            operators = ['>=', '<=', '==', '!=', '>', '<']
            for op in operators:
                if condition.startswith(op):
                    try:
                        value = float(condition[len(op):].strip())
                        if op == '>=':
                            return self.data[column] >= value
                        elif op == '<=':
                            return self.data[column] <= value
                        elif op == '==':
                            return self.data[column] == value
                        elif op == '!=':
                            return self.data[column] != value
                        elif op == '>':
                            return self.data[column] > value
                        elif op == '<':
                            return self.data[column] < value
                    except ValueError:
                        pass
            
            # 如果没有运算符，默认为等于
            try:
                value = float(condition)
                return self.data[column] == value
            except ValueError:
                print(f"错误: 无法解析数值条件 '{condition}'")
                return pd.Series([True] * len(self.data))
                
        except Exception as e:
            print(f"解析数值条件失败: {e}")
            return pd.Series([True] * len(self.data))
    
    def _is_null_condition(self, condition):
        """检查是否为空值条件"""
        condition = condition.strip().lower()
        null_keywords = ['null', 'nan', 'empty', 'isnull', 'isna', 'notnull', 'notna', 
                        '空值', '空', '非空', '不为空', '为空']
        return any(keyword in condition for keyword in null_keywords)
    
    def _parse_null_condition(self, column, condition):
        """解析空值条件"""
        try:
            condition = condition.strip().lower()
            column_data = self.data[column]
            
            # 支持的空值条件格式:
            # null, isnull, isna, 空值, 为空 - 查找空值
            # notnull, notna, 非空, 不为空 - 查找非空值  
            # empty - 查找空字符串
            # notempty - 查找非空字符串
            
            if condition in ['null', 'isnull', 'isna', '空值', '为空']:
                # 查找空值 (包括 NaN, None, pd.NA)
                return column_data.isnull()
            
            elif condition in ['notnull', 'notna', '非空', '不为空']:
                # 查找非空值
                return column_data.notnull()
            
            elif condition in ['empty', '空']:
                # 查找空字符串 (对于字符串列)
                if pd.api.types.is_string_dtype(column_data) or column_data.dtype == 'object':
                    # 同时检查空值和空字符串
                    return column_data.isnull() | (column_data.astype(str).str.strip() == '')
                else:
                    # 对于非字符串列，只检查空值
                    return column_data.isnull()
            
            elif condition in ['notempty', '非空字符串']:
                # 查找非空字符串
                if pd.api.types.is_string_dtype(column_data) or column_data.dtype == 'object':
                    # 既不是空值也不是空字符串
                    return column_data.notnull() & (column_data.astype(str).str.strip() != '')
                else:
                    # 对于非字符串列，只检查非空值
                    return column_data.notnull()
            
            else:
                print(f"警告: 未知的空值条件 '{condition}'，支持的条件: null/isnull/isna/空值/为空 (查找空值), notnull/notna/非空/不为空 (查找非空值), empty/空 (查找空字符串), notempty/非空字符串 (查找非空字符串)")
                return pd.Series([True] * len(self.data))
                
        except Exception as e:
            print(f"解析空值条件失败: {e}")
            return pd.Series([True] * len(self.data))
    
    def _parse_string_condition(self, column, condition):
        """解析字符串型条件"""
        try:
            condition = condition.strip()
            
            # 支持的字符串条件格式:
            # exact:value (精确匹配)
            # contains:value (包含)
            # startswith:value (开头匹配)
            # endswith:value (结尾匹配)
            # regex:pattern (正则表达式)
            # in:value1,value2,value3 (多个值匹配)
            # value1,value2,value3 (多个值匹配的简化语法)
            # !=value (不等于)
            # 默认为精确匹配
            
            # 检查是否为不等于条件
            if condition.startswith('!='):
                value = condition[2:].strip()
                print(f"检测到不等于条件: {value}")
                return self.data[column].astype(str) != value
            
            # 检查是否为多值匹配的简化语法 (逗号分隔)
            if ',' in condition and ':' not in condition:
                values = [v.strip() for v in condition.split(',')]
                print(f"检测到多值匹配: {values}")
                return self.data[column].astype(str).isin(values)
            
            if ':' in condition:
                method, value = condition.split(':', 1)
                method = method.lower().strip()
                value = value.strip()
                
                if method == 'exact':
                    return self.data[column].astype(str) == value
                elif method == 'contains':
                    return self.data[column].astype(str).str.contains(value, case=False, na=False)
                elif method == 'startswith':
                    return self.data[column].astype(str).str.startswith(value, na=False)
                elif method == 'endswith':
                    return self.data[column].astype(str).str.endswith(value, na=False)
                elif method == 'regex':
                    return self.data[column].astype(str).str.contains(value, regex=True, case=False, na=False)
                elif method == 'in':
                    values = [v.strip() for v in value.split(',')]
                    return self.data[column].astype(str).isin(values)
                else:
                    print(f"警告: 未知的字符串匹配方法 '{method}'，使用精确匹配")
                    return self.data[column].astype(str) == value
            else:
                # 默认精确匹配
                return self.data[column].astype(str) == condition
                
        except Exception as e:
            print(f"解析字符串条件失败: {e}")
            return pd.Series([True] * len(self.data))
    
    def apply_filters(self):
        """应用所有过滤条件"""
        if not self.filters:
            print("警告: 没有设置任何过滤条件，返回原始数据")
            return self.data
        
        try:
            # 初始化掩码为全True
            mask = pd.Series([True] * len(self.data))
            
            # 应用每个过滤条件 (AND逻辑)
            for filter_item in self.filters:
                column = filter_item['column']
                condition = filter_item['condition']
                
                print(f"应用过滤条件: {column} -> {condition}")
                condition_mask = self.parse_condition(column, condition)
                mask = mask & condition_mask
                
                print(f"过滤后剩余行数: {mask.sum()}")
            
            filtered_data = self.data[mask]
            print(f"\n最终过滤结果: {len(filtered_data)} 行 (原始: {len(self.data)} 行)")
            
            return filtered_data
            
        except Exception as e:
            print(f"应用过滤条件失败: {e}")
            return self.data
    
    def save_results(self, filtered_data):
        """保存过滤结果"""
        try:
            if self.output_file.suffix.lower() == '.csv':
                filtered_data.to_csv(self.output_file, index=False, encoding='utf-8')
            elif self.output_file.suffix.lower() in ['.xls', '.xlsx']:
                filtered_data.to_excel(self.output_file, index=False)
            else:
                print(f"错误: 不支持的输出格式 {self.output_file.suffix}")
                return False
            
            print(f"过滤结果已保存到: {self.output_file}")
            return True
            
        except Exception as e:
            print(f"保存结果失败: {e}")
            return False
    
    def process(self, filter_conditions):
        """执行完整的过滤流程"""
        print("=" * 50)
        print("🔍 CSV矩阵过滤工具")
        print("=" * 50)
        
        # 加载数据
        if not self.load_data():
            return False
        
        # 添加过滤条件
        for condition in filter_conditions:
            # 使用冒号作为分隔符：column:condition
            if ':' not in condition:
                print(f"错误: 过滤条件格式错误 '{condition}'，应为 'column:condition'")
                return False
            
            column, cond = condition.split(':', 1)
            column = column.strip()
            cond = cond.strip()
            
            if not self.add_filter(column, cond):
                return False
        
        # 应用过滤条件
        filtered_data = self.apply_filters()
        
        if len(filtered_data) == 0:
            print("警告: 过滤后没有数据行")
        
        # 保存结果
        if not self.save_results(filtered_data):
            return False
        
        # 显示摘要信息
        print("\n📋 过滤摘要:")
        print(f"📁 输入文件: {self.input_file}")
        print(f"📄 输出文件: {self.output_file}")
        print(f"🔍 过滤条件: {len(self.filters)} 个")
        for i, filter_item in enumerate(self.filters, 1):
            print(f"  {i}. {filter_item['column']} -> {filter_item['condition']}")
        print(f"📊 原始行数: {len(self.data)}")
        print(f"📈 过滤后行数: {len(filtered_data)}")
        print(f"📉 过滤比例: {len(filtered_data)/len(self.data)*100:.1f}%")
        
        return True

def main():
    parser = argparse.ArgumentParser(
        description='CSV/Excel矩阵过滤工具',
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
使用示例:
  # 数值型过滤
  matrix-filter input.csv -o output.csv --filter "age:>=18" "score:>80"
  matrix-filter input.xlsx -o output.xlsx --filter "price:>100"  # Excel文件过滤
  matrix-filter input.csv -o output.csv --filter "price:10-50"  # 范围过滤
  matrix-filter input.csv -o output.csv --filter "category:1,2,3"  # 多值过滤
  
  # 字符串型过滤
  matrix-filter input.csv -o output.csv --filter "name:exact:John"  # 精确匹配
  matrix-filter input.csv -o output.csv --filter "city:contains:北京"  # 包含匹配
  matrix-filter input.csv -o output.csv --filter "email:endswith:.com"  # 结尾匹配
  matrix-filter input.csv -o output.csv --filter "phone:regex:^1[3-9]\\d{9}$"  # 正则匹配
  matrix-filter input.csv -o output.csv --filter "status:in:active,pending"  # 多值匹配
  matrix-filter input.csv -o output.csv --filter "batch:第二批,第三批"  # 多值匹配简化语法
  
  # 空值过滤
  matrix-filter input.csv -o output.csv --filter "phone:null"  # 提取电话号码为空的行
  matrix-filter input.csv -o output.csv --filter "email:空值"  # 提取邮箱为空的行
  matrix-filter input.csv -o output.csv --filter "address:empty"  # 提取地址为空字符串的行
  matrix-filter input.csv -o output.csv --filter "name:notnull"  # 提取姓名不为空的行
  matrix-filter input.csv -o output.csv --filter "description:非空"  # 提取描述不为空的行
  
  # 组合过滤 (AND逻辑)
  matrix-filter input.csv -o output.csv --filter "age:>=18" "city:contains:上海" "score:>85"
  matrix-filter input.csv -o output.csv --filter "phone:notnull" "email:notnull"  # 电话和邮箱都不为空
  
  # 批次过滤示例
  matrix-filter input.csv -o output.csv --filter "批次:第二批,第三批"  # 获取第二批和第三批的数据
  matrix-filter input.csv -o output.csv --filter "status:活跃,待审核,已完成"  # 获取多个状态的数据

过滤条件格式:
  数值型: >10, >=10, <10, <=10, ==10, !=10, 10-20, 10,20,30
  字符串型: exact:value, contains:value, startswith:value, endswith:value, regex:pattern, in:value1,value2
           或直接使用逗号分隔的多值: value1,value2,value3
  空值型: null/isnull/isna/空值/为空 (查找空值), notnull/notna/非空/不为空 (查找非空值)
         empty/空 (查找空字符串), notempty/非空字符串 (查找非空字符串)
        """
    )
    parser.add_argument('input_file', help='输入CSV/Excel文件路径')
    parser.add_argument('-o', '--output', required=True, help='输出CSV/Excel文件路径')
    parser.add_argument('--filter', nargs='+', required=True, 
                       help='过滤条件，格式为 "column:condition"，多个条件用空格分隔')
    
    args = parser.parse_args()
    
    filter_tool = MatrixFilter(args.input_file, args.output)
    success = filter_tool.process(args.filter)
    
    if success:
        print("\n🎉 CSV矩阵过滤工具执行成功!")
    else:
        print("\n❌ CSV矩阵过滤工具执行失败!")
        exit(1)

if __name__ == '__main__':
    main()